import numpy as np
import matplotlib.pyplot as plt
from typing import List,Tuple,Union,Dict

def calNormal(n1:np.ndarray,n2:np.ndarray,n3:np.ndarray)->np.ndarray:
    """,according to right-hand rule; cal triangle's normal vector,by given three nodes coords,"""
    # define vector
    v12,v13=n2-n1,n3-n1
    # cal the cross
    normal=np.cross(v12,v13)
    n=normal/np.linalg.norm(normal)
    return n

def imShowMat(mat:np.ndarray,title:str)->None:
    # 使用 imshow 可视化矩阵
    plt.imshow(mat, cmap='jet', interpolation='nearest')  # cmap 是颜色映射
    plt.colorbar()  # 显示颜色条
    plt.title(title)
    plt.show()

def DetMatrixByLu(A:np.ndarray)->np.ndarray:
    """
    利用LU分解求矩阵的行列式
    :param A: 方阵
    :return: 矩阵的行列式
    """
    from scipy.linalg import lu
    # 进行 LU 分解
    P, L, U = lu(A)
    # 计算行列式
    det_L = np.prod(np.diag(L))
    det_U = np.prod(np.diag(U))
    det_A = det_L * det_U * np.linalg.det(P)
    return det_A

def initNumpyPrintOptions():
    """
    Initialize the numpy print options to the default values.
    """
    np.set_printoptions(linewidth=250, precision=8, suppress=True)

def rotate_90(vector, clockwise=False):
    """
    Rotate a 2D vector by 90 degrees around the z-axis.
    :param vector: Input vector (a_x, a_y)
    :param clockwise: True-顺时针旋转, False-逆时针旋转
    :return: Rotated vector
    """
    a_x, a_y = vector
    if clockwise:
        return np.array([a_y, -a_x])  # 顺时针旋转
    else:
        return np.array([-a_y, a_x])  # 逆时针旋转
    
def dofInds(NodeLabel:int=1,nodeDofNum:int=1)->List[int]:
    """给定节点编号i和自由度个数n, 计算节点i的所有自由度的索引(0-based) """
    i,n=int(NodeLabel),int(nodeDofNum)
    return (np.array(list(range(n)))+(i-1)*n).tolist()

def get_ithdof_index(NodeLabel:int=1,dofNum:int=1,ithdof:int=1)->int:
    """给定节点编号i,自由度个数n, 返回节点中第j个自由度的全局索引(0-based) """
    i,n,j=NodeLabel,dofNum,ithdof
    return dofInds(i,n)[j-1]

def BCprocess(K:np.ndarray,F:np.ndarray,bc_dofs:List[int],bc_values:List[float])->Tuple[np.ndarray,np.ndarray]:
    """处理边界条件,直大数法
    Args:
        K: 刚度矩阵
        F: 外力列向量
        bc_dofs: 边界条件自由度的全局编号list
        bc_values: 边界条件值list
    Returns: 
        K: 处理后的刚度矩阵
        F: 处理后的外力向量
    """
    if len(bc_dofs)!=len(bc_values):
        raise ValueError("bc_dofs and bc_values must have the same length")
    # 处理边界条件
    K_=K.copy()
    F_=F.copy()
    C=np.max(np.abs(K))*1e8
    for dof,value in zip(bc_dofs,bc_values):
        K_[dof,dof]+=C
        F_[dof,0]+=value*C
    return K_,F_

def export_csv(matrix:np.ndarray,
                output_file:str,
                norm_array:bool=False,
                )->bool:
    """将矩阵导出为csv文件
    """
    import numpy as np
    # 假设你有一个 NumPy 数组
    if matrix.dtype==complex:
        array = np.array(matrix,dtype=complex)
    else:
        array = np.array(matrix,dtype=float)
        
    # 普通的数值矩阵,直接导出
    if norm_array:
    # 使用 savetxt 函数将数组保存为 CSV 文件,保留 6 位有效数字,科学计数法
        np.savetxt(output_file, array, delimiter=',', fmt='%.6e')
        return True
    else:
        # 使用条件表达式,选择file mode
        with open(output_file, mode = 'w',newline='\n') as f:
            
            for index,value in np.ndenumerate(array):
                # 二维数组
                if len(index)==2:
                    # 对于正数和负数,保留不同的有效数字位数
                    f.write(f"{value:.9e}") if value<0 else f.write(f"{value:.10e}")
                    # 判断是否是该行最后一个元素,如果不是,则添加一个逗号和空格
                    f.write(f"\n") if index[1]==array.shape[1]-1 else f.write(" , ")
                # 一维数组
                elif len(index)==1: 
                    f.write(f"{value:.5e}") if value>0 else  f.write(f"{value:.6e}")
                    f.write(f"\n") if index[0]==array.shape[0]-1 else f.write(" , ")
        return True

def get_entry(line):
    """将质量/刚度矩阵文件的 一行数据转换为矩阵的元素
    """
    data=line.strip().split()
    if len(data)==3:
        # FORMAT="coordinate"
        entry=[float(data[0]),float(data[1]),float(data[2])]
    if len(data)==5:
        # FORMAT="label"
        pass
    return entry # type: ignore
    
def read_mtx(file_name:str):
    with open(file=file_name,mode='r') as f:
        datas_list=[]
        while True:
            line=f.readline()
            
            if not line:
                break
            entry=get_entry(line)
            datas_list.append(entry)
            
        # dt = np.dtype([('row', int),('col', int),('value', float)]) 
        # datas=np.array(datas_list,dtype=dt)
        datas=np.array(datas_list)
        matrix=np.zeros((int(np.max(datas[:,0])),int(np.max(datas[:,1]))))
        for _,i in enumerate(datas):
            row=int(i[0])-1
            col=int(i[1])-1
            value=i[2]
            matrix[row,col]=value
    return matrix

def plotMatrix(matrix:np.ndarray):
    import numpy as np
    import matplotlib.pyplot as plt

    # 使用 imshow 可视化矩阵
    plt.imshow(matrix, cmap='viridis', interpolation='nearest')  # 可以选择不同的颜色映射(cmap)
    plt.colorbar()  # 添加颜色条
    plt.title("Matrix Visualization")
    plt.show()
    
def plot1d(array:np.ndarray,**kwargs)->None:
    import matplotlib.pyplot as plt
    if 'title' not in kwargs.keys():
        kwargs['title'] = ''
    # 绘制数组
    plt.plot(array, marker='o')

    # 添加标题和标签
    plt.title(f"1D Array Plot\n{kwargs['title']}")
    plt.xlabel('Array Index')
    plt.ylabel('Array Value')
    
    # y刻度分成8分
    if 'showData' in kwargs.keys() and kwargs['showData']==True:
        ytick_interval=(np.max(array)*1.3-np.min(array)*0.8)/8.0 
        plt.yticks(np.arange(np.min(array)*0.8, np.max(array)*1.3,ytick_interval))
        for x, y in enumerate(array):
            plt.text(x-0.3, y+0.1*ytick_interval, f"{y:.2e}", ha='center', va='bottom', fontsize=9)

    # 显示图形
    plt.show()

def dofGen(s:int,e:int)->np.ndarray:
    dofs=[]
    for i in range(s,e+1):
        dofs.append(i*3-2)
        dofs.append(i*3-1)
        dofs.append(i*3)
    return np.array(dofs)-1

def read_table(f)->list:
    """read table data in file"""
    data=[]
    while True:
        line=f.readline()
        p=line.strip('\n').strip(' ')
        # 遇到空行就退出
        if not p:
            break
        else:
            print(line)
            line=line.strip().split()
            entry=list(map(float,line[1:]))
            data+=entry
    
    return data

def read_U_from_rpt(file_name:str,start_lines:list=[1,],**kwargs):
    """read data from *.rpt file"""
    eig_vec=[]
    
    with open(file_name, 'r') as f:
        for index, line in enumerate(f,1):
            if index in start_lines:
                eig_vec.append(read_table(f))
            
    eig_vec=np.array(eig_vec)
    return eig_vec
from typing import List,Tuple,Dict

def export_vtk(output_file:str,
                nodes:np.ndarray,
                element_Blocks:List[np.ndarray],
                element_types:List[str],
                point_data= {},
                cell_data= {},
                **kwargs)->bool:
    """导出vtk格式的网格数据文件
    inputs:
    output_file: 输出文件名
    nodes: 节点坐标矩阵,shape=(n,3)
    element_Blocks: 单元列表,
    element_types: 单元类型列表,
    point_data: 节点数据字典.e.g. {'SCALARS':{'disp':disp_array,'stress':stress_array},'VECTORS':{'velocity':velocity_array}}
    cell_data: 单元数据字典.e.g. {'SCALARS':{'disp':disp_array,'stress':stress_array},'VECTORS':{'velocity':velocity_array}}
    """
    def cal_cell_num(element_Blocks:List[np.ndarray])->tuple:
        cell_num,cell_piont_num=0,0
        for _,block in enumerate(element_Blocks):
            cell_num+=block.shape[0]
            cell_piont_num+=block.shape[0]*(block.shape[1]+1)
        return cell_num,cell_piont_num
    
    SEP="  " if 'sep'not in kwargs.keys()  else kwargs['sep']
    
    type_dict={"beam": 3, # vtk_line, 梁单元
               "tri": 5, # vtk_triangle, 三角形单元
               "quad":9, # vtk_quad, 四边形单元
               "tet":10, # vtk_tetra, 四面体单元
               "hex":12, # vtk_hexahedron, 六面体单元
               "wedge":13, # vtk_wedge, 三棱柱单元
               "pyramid":14, # vtk_pyramid, 四棱锥单元
               }
    
    f=open(output_file,'w')
    
    # 写入文件头
    f.write('# vtk DataFile Version 3.0\n')
    f.write(f'{kwargs["title"]}\n') if 'title' in kwargs.keys() else f.write('Finite Element Results\n')
    f.write('ASCII\n')
    f.write('DATASET UNSTRUCTURED_GRID\n')
    
    # 写入点坐标
    f.write(f'\nPOINTS {nodes.shape[0]} float\n')
    for i in range(nodes.shape[0]):
        f.write(SEP.join(map(str,nodes[i,:]))+'\n')
        
    # 写入单元数据
    n,size=cal_cell_num(element_Blocks)
    
    # 构造单元数据
    f.write(f'\nCELLS {n} {size}\n')
    cells_lines,cell_type_lines="",""
    for block ,typeStr in zip(element_Blocks,element_types):
        for i in range(block.shape[0]):
            cells_lines+=f"{block.shape[1]}"+" " +SEP.join(map(str,block[i,:]))+'\n'
            cell_type_lines+=f"{type_dict[typeStr]}\n"
    f.write(cells_lines)
    
    # 写入单元类型数据
    f.write(f'\nCELL_TYPES {n}\n')
    f.write(cell_type_lines)
    
    # 写入节点数据,节点字典数据为空则跳过
    if point_data:
        f.write(f'\nPOINT_DATA {nodes.shape[0]}\n') 
        if 'SCALARS' in point_data.keys():
            for var_name,array in point_data['SCALARS'].items():
                if array.ndim==1:
                    # 如果是一维数组,则扩展成二维数组
                    array=array.reshape((-1,1))
                # 标量的Ncomp=1~4
                f.write(f'SCALARS {var_name} float {array.shape[1]}\n')
                f.write(f'LOOKUP_TABLE default\n')
                for i in range(array.shape[0]):
                    f.write(SEP.join(map(str,array[i,:]))+'\n')
        
        # VECTORS dataName dataType
        # v0x v0y v0z
        # v1x v1y v1z
        # …
        # v(n-1)x v(n-1)y v(n-1)z
        if 'VECTORS' in point_data.keys():
            for var_name,array in point_data['VECTORS'].items():
                if array.ndim==1:
                    # 如果是一维数组,则扩展成二维数组,向量必须是(-1,3)
                    array=array.reshape((-1,3))
                assert array.shape[1]==3, "VECTORS data must be (-1,3) shape"
                
                f.write(f'VECTORS {var_name} float\n')
                for i in range(array.shape[0]):
                    f.write(SEP.join(map(str,array[i,:]))+'\n')
    else:
        pass
    # 写入单元数据,单元字典数据为空则跳过
    if cell_data:
        f.write(f'\nCELL_DATA {n}\n')
        if 'SCALARS' in cell_data.keys():
            for var_name,array in cell_data['SCALARS'].items():
                if array.ndim==1:
                    # 如果是一维数组,则扩展成二维数组
                    array=array.reshape((-1,1))
                # 标量的Ncomp=1~4
                f.write(f'SCALARS {var_name} float {array.shape[1]}\n')
                f.write(f'LOOKUP_TABLE default\n')
                for i in range(array.shape[0]):
                    f.write(SEP.join(map(str,array[i,:]))+'\n')
                    
        if 'VECTORS' in cell_data.keys():
            for var_name,array in cell_data['VECTORS'].items():
                if array.ndim==1:
                    # 如果是一维数组,则扩展成二维数组,向量必须是(-1,3)
                    array=array.reshape((-1,3))
                assert array.shape[1]==3, "VECTORS data must be (-1,3) shape"
                
                f.write(f'VECTORS {var_name} float\n')
                for i in range(array.shape[0]):
                    f.write(SEP.join(map(str,array[i,:]))+'\n')
    else:
        pass
    f.close()
    return True 
